package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"time"

	"adkr/pkg/crypto"
)

type BenchmarkResult struct {
	CurveType    string        `json:"curve_type"`
	Operation    string        `json:"operation"`
	Duration     time.Duration `json:"duration_ms"`
	Iterations   int           `json:"iterations"`
	AvgPerOp     float64       `json:"avg_per_op_ms"`
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: curve_bench <curve_type>")
		fmt.Println("  curve_type: bls12-381 or secp256k1")
		os.Exit(1)
	}
	
	curveType := os.Args[1]
	
	var results []BenchmarkResult
	
	switch curveType {
	case "bls12-381":
		crypto.SetCurveType(crypto.BLS12_381)
		results = benchmarkBLS12381()
	case "secp256k1":
		crypto.SetCurveType(crypto.Secp256k1)
		results = benchmarkSecp256k1()
	default:
		log.Fatalf("Unknown curve type: %s", curveType)
	}
	
	// Output results as JSON
	jsonData, err := json.MarshalIndent(results, "", "  ")
	if err != nil {
		log.Fatalf("Failed to marshal results: %v", err)
	}
	
	fmt.Println(string(jsonData))
}

func benchmarkBLS12381() []BenchmarkResult {
	var results []BenchmarkResult
	
	// Generate threshold keys
	start := time.Now()
	_, pks, sks := crypto.GenerateThresholdKeys(4, 3)
	duration := time.Since(start)
	
	results = append(results, BenchmarkResult{
		CurveType:  "bls12-381",
		Operation:  "generate_threshold_keys",
		Duration:   duration,
		Iterations: 1,
		AvgPerOp:   float64(duration.Nanoseconds()) / 1e6,
	})
	
	// Sign message
	message := []byte("test message for benchmarking")
	iterations := 100
	
	start = time.Now()
	for i := 0; i < iterations; i++ {
		_ = crypto.Sign(&sks[0], message)
	}
	duration = time.Since(start)
	
	results = append(results, BenchmarkResult{
		CurveType:  "bls12-381",
		Operation:  "sign",
		Duration:   duration,
		Iterations: iterations,
		AvgPerOp:   float64(duration.Nanoseconds()) / float64(iterations) / 1e6,
	})
	
	// Verify signature
	sig := crypto.Sign(&sks[0], message)
	
	start = time.Now()
	for i := 0; i < iterations; i++ {
		_ = crypto.VerifyShare(&pks[0], sig, message)
	}
	duration = time.Since(start)
	
	results = append(results, BenchmarkResult{
		CurveType:  "bls12-381",
		Operation:  "verify_share",
		Duration:   duration,
		Iterations: iterations,
		AvgPerOp:   float64(duration.Nanoseconds()) / float64(iterations) / 1e6,
	})
	
	return results
}

func benchmarkSecp256k1() []BenchmarkResult {
	var results []BenchmarkResult
	
	// Generate threshold keys
	start := time.Now()
	keySet := crypto.GenerateSecp256k1ThresholdKeys(4, 3)
	duration := time.Since(start)
	
	results = append(results, BenchmarkResult{
		CurveType:  "secp256k1",
		Operation:  "generate_threshold_keys",
		Duration:   duration,
		Iterations: 1,
		AvgPerOp:   float64(duration.Nanoseconds()) / 1e6,
	})
	
	// Sign message
	message := []byte("test message for benchmarking")
	iterations := 100
	
	start = time.Now()
	for i := 0; i < iterations; i++ {
		_ = crypto.SignSecp256k1Share(keySet.Shares[0].Share, message)
	}
	duration = time.Since(start)
	
	results = append(results, BenchmarkResult{
		CurveType:  "secp256k1",
		Operation:  "sign",
		Duration:   duration,
		Iterations: iterations,
		AvgPerOp:   float64(duration.Nanoseconds()) / float64(iterations) / 1e6,
	})
	
	// Verify signature
	sig := crypto.SignSecp256k1Share(keySet.Shares[0].Share, message)
	
	start = time.Now()
	for i := 0; i < iterations; i++ {
		_ = crypto.VerifySecp256k1Share(keySet.Shares[0].PubKey, sig, message)
	}
	duration = time.Since(start)
	
	results = append(results, BenchmarkResult{
		CurveType:  "secp256k1",
		Operation:  "verify_share",
		Duration:   duration,
		Iterations: iterations,
		AvgPerOp:   float64(duration.Nanoseconds()) / float64(iterations) / 1e6,
	})
	
	return results
}
